"""
23_graph_rag_basics.py
======================
Graph RAG (Retrieval-Augmented Generation) Fundamentals

This module covers:
1. Understanding Graph RAG architecture
2. Building a Graph RAG pipeline
3. Combining vector search with graph traversal
4. Context assembly for LLM
5. End-to-end RAG implementation

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional, Callable
from dataclasses import dataclass, field
from abc import ABC, abstractmethod
import json
from datetime import datetime
from neo4j import Driver
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


@dataclass
class RetrievedContext:
    """Context retrieved from the knowledge graph."""
    chunks: List[Dict[str, Any]]
    entities: List[Dict[str, Any]]
    relationships: List[Dict[str, Any]]
    graph_context: str
    total_tokens: int = 0
    
    def to_prompt_context(self) -> str:
        """Format context for LLM prompt."""
        sections = []
        
        # Add chunk content
        if self.chunks:
            sections.append("=== RELEVANT DOCUMENT SECTIONS ===")
            for i, chunk in enumerate(self.chunks, 1):
                sections.append(f"\n[Section {i}] (Relevance: {chunk.get('score', 0):.2f})")
                sections.append(chunk.get('content', ''))
        
        # Add entity information
        if self.entities:
            sections.append("\n=== RELATED ENTITIES ===")
            for entity in self.entities:
                sections.append(f"- {entity.get('label', 'Entity')}: {entity.get('name', 'Unknown')}")
                if entity.get('properties'):
                    for k, v in entity['properties'].items():
                        sections.append(f"    {k}: {v}")
        
        # Add relationship context
        if self.graph_context:
            sections.append("\n=== KNOWLEDGE GRAPH CONTEXT ===")
            sections.append(self.graph_context)
        
        return "\n".join(sections)


@dataclass
class RAGQuery:
    """A RAG query with configuration."""
    question: str
    top_k: int = 5
    include_graph_context: bool = True
    max_context_tokens: int = 4000
    min_similarity: float = 0.3
    context_hops: int = 1


@dataclass
class RAGResponse:
    """Response from RAG system."""
    question: str
    answer: str
    context: RetrievedContext
    sources: List[str]
    confidence: float
    processing_time_ms: float


class EmbeddingProvider(ABC):
    """Abstract base class for embedding providers."""
    
    @abstractmethod
    def embed(self, text: str) -> List[float]:
        """Generate embedding for text."""
        pass
    
    @abstractmethod
    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for multiple texts."""
        pass


class OpenAIEmbeddings(EmbeddingProvider):
    """OpenAI embeddings provider."""
    
    def __init__(self, model: str = "text-embedding-3-small"):
        self.model = model
        self._client = None
    
    @property
    def client(self):
        if self._client is None:
            try:
                from openai import OpenAI
                self._client = OpenAI()
            except ImportError:
                raise ImportError("openai package required")
        return self._client
    
    def embed(self, text: str) -> List[float]:
        response = self.client.embeddings.create(
            model=self.model,
            input=text
        )
        return response.data[0].embedding
    
    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        response = self.client.embeddings.create(
            model=self.model,
            input=texts
        )
        return [d.embedding for d in response.data]


class LocalEmbeddings(EmbeddingProvider):
    """Local embeddings using sentence-transformers."""
    
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model_name = model_name
        self._model = None
    
    @property
    def model(self):
        if self._model is None:
            try:
                from sentence_transformers import SentenceTransformer
                self._model = SentenceTransformer(self.model_name)
            except ImportError:
                raise ImportError("sentence-transformers package required")
        return self._model
    
    def embed(self, text: str) -> List[float]:
        return self.model.encode(text).tolist()
    
    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        embeddings = self.model.encode(texts)
        return [e.tolist() for e in embeddings]


class GraphRAG:
    """
    Graph RAG implementation combining vector search with knowledge graphs.
    
    Architecture:
    
    ┌─────────────────────────────────────────────────────────────────┐
    │                      USER QUERY                                 │
    └───────────────────────────┬─────────────────────────────────────┘
                                │
                                ▼
    ┌─────────────────────────────────────────────────────────────────┐
    │                    QUERY PROCESSING                             │
    │  ┌──────────────┐  ┌──────────────┐  ┌──────────────┐          │
    │  │   Embed      │  │   Extract    │  │   Identify   │          │
    │  │   Query      │  │   Entities   │  │   Intent     │          │
    │  └──────────────┘  └──────────────┘  └──────────────┘          │
    └───────────────────────────┬─────────────────────────────────────┘
                                │
                                ▼
    ┌─────────────────────────────────────────────────────────────────┐
    │                    RETRIEVAL STAGE                              │
    │  ┌──────────────┐  ┌──────────────┐  ┌──────────────┐          │
    │  │   Vector     │  │   Graph      │  │   Hybrid     │          │
    │  │   Search     │  │   Traversal  │  │   Search     │          │
    │  │   (HNSW)     │  │   (Cypher)   │  │   (Both)     │          │
    │  └──────────────┘  └──────────────┘  └──────────────┘          │
    └───────────────────────────┬─────────────────────────────────────┘
                                │
                                ▼
    ┌─────────────────────────────────────────────────────────────────┐
    │                    CONTEXT ASSEMBLY                             │
    │  ┌──────────────┐  ┌──────────────┐  ┌──────────────┐          │
    │  │   Rank &     │  │   Expand     │  │   Format     │          │
    │  │   Filter     │  │   Context    │  │   for LLM    │          │
    │  └──────────────┘  └──────────────┘  └──────────────┘          │
    └───────────────────────────┬─────────────────────────────────────┘
                                │
                                ▼
    ┌─────────────────────────────────────────────────────────────────┐
    │                    GENERATION STAGE                             │
    │  ┌────────────────────────────────────────────────────────┐    │
    │  │              LLM (GPT-4, Claude, etc.)                 │    │
    │  │  Prompt = Instructions + Context + Query               │    │
    │  └────────────────────────────────────────────────────────┘    │
    └───────────────────────────┬─────────────────────────────────────┘
                                │
                                ▼
    ┌─────────────────────────────────────────────────────────────────┐
    │                      RESPONSE                                   │
    └─────────────────────────────────────────────────────────────────┘
    
    Key Advantages of Graph RAG:
    1. Semantic search finds relevant chunks
    2. Graph traversal provides related entities
    3. Relationships add structured context
    4. Path information explains connections
    """
    
    def __init__(
        self,
        driver: Driver,
        embedding_provider: EmbeddingProvider = None,
        vector_index_name: str = "chunk_embeddings"
    ):
        self.driver = driver
        self.embedding_provider = embedding_provider or LocalEmbeddings()
        self.vector_index_name = vector_index_name
    
    # =========================================================
    # RETRIEVAL METHODS
    # =========================================================
    
    def vector_search(
        self,
        query_embedding: List[float],
        top_k: int = 5,
        min_score: float = 0.0
    ) -> List[Dict[str, Any]]:
        """
        Perform vector similarity search.
        
        Args:
            query_embedding: Query vector
            top_k: Number of results
            min_score: Minimum similarity threshold
            
        Returns:
            Similar chunks with scores
        """
        query = """
        CALL db.index.vector.queryNodes($index_name, $k, $embedding)
        YIELD node, score
        WHERE score >= $min_score
        RETURN 
            node.id AS id,
            node.content AS content,
            node.index AS chunk_index,
            score
        ORDER BY score DESC
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                index_name=self.vector_index_name,
                k=top_k,
                embedding=query_embedding,
                min_score=min_score
            )
            return [record.data() for record in result]
    
    def get_chunk_entities(
        self,
        chunk_ids: List[str]
    ) -> List[Dict[str, Any]]:
        """
        Get entities related to chunks.
        
        Args:
            chunk_ids: List of chunk IDs
            
        Returns:
            Related entities
        """
        query = """
        UNWIND $chunk_ids AS cid
        MATCH (c:Chunk {id: cid})
        OPTIONAL MATCH (c)-[:MENTIONS|RELATES_TO|ABOUT]->(e)
        WHERE e IS NOT NULL
        RETURN DISTINCT
            labels(e)[0] AS label,
            e.name AS name,
            e.id AS id,
            properties(e) AS properties
        """
        
        with self.driver.session() as session:
            result = session.run(query, chunk_ids=chunk_ids)
            return [record.data() for record in result]
    
    def get_entity_context(
        self,
        entity_ids: List[str],
        hops: int = 1
    ) -> str:
        """
        Get graph context around entities.
        
        Args:
            entity_ids: Entity IDs to explore
            hops: Number of relationship hops
            
        Returns:
            Formatted graph context
        """
        query = f"""
        UNWIND $entity_ids AS eid
        MATCH (e {{id: eid}})
        OPTIONAL MATCH path = (e)-[r*1..{hops}]-(related)
        WHERE related IS NOT NULL
        WITH e, related, r
        RETURN DISTINCT
            labels(e)[0] AS source_label,
            e.name AS source_name,
            [rel IN r | type(rel)][0] AS relationship,
            labels(related)[0] AS target_label,
            related.name AS target_name
        LIMIT 50
        """
        
        with self.driver.session() as session:
            result = session.run(query, entity_ids=entity_ids)
            records = [record.data() for record in result]
            
            # Format as natural language
            context_lines = []
            for r in records:
                if r['relationship'] and r['target_name']:
                    line = f"- {r['source_name']} ({r['source_label']}) " \
                           f"--[{r['relationship']}]--> " \
                           f"{r['target_name']} ({r['target_label']})"
                    context_lines.append(line)
            
            return "\n".join(context_lines)
    
    def hybrid_retrieve(
        self,
        query: RAGQuery
    ) -> RetrievedContext:
        """
        Perform hybrid retrieval combining vector search and graph traversal.
        
        Args:
            query: RAG query configuration
            
        Returns:
            Retrieved context
        """
        import time
        start = time.perf_counter()
        
        # 1. Embed query
        query_embedding = self.embedding_provider.embed(query.question)
        
        # 2. Vector search
        chunks = self.vector_search(
            query_embedding,
            top_k=query.top_k,
            min_score=query.min_similarity
        )
        
        # 3. Get related entities
        chunk_ids = [c['id'] for c in chunks if c.get('id')]
        entities = self.get_chunk_entities(chunk_ids) if chunk_ids else []
        
        # 4. Get graph context
        graph_context = ""
        if query.include_graph_context and entities:
            entity_ids = [e['id'] for e in entities if e.get('id')]
            if entity_ids:
                graph_context = self.get_entity_context(
                    entity_ids,
                    hops=query.context_hops
                )
        
        # 5. Get relationships between entities
        relationships = self._get_inter_entity_relationships(
            [e['id'] for e in entities if e.get('id')]
        )
        
        return RetrievedContext(
            chunks=chunks,
            entities=entities,
            relationships=relationships,
            graph_context=graph_context
        )
    
    def _get_inter_entity_relationships(
        self,
        entity_ids: List[str]
    ) -> List[Dict[str, Any]]:
        """Get relationships between retrieved entities."""
        if len(entity_ids) < 2:
            return []
        
        query = """
        UNWIND $ids AS id1
        UNWIND $ids AS id2
        WITH id1, id2 WHERE id1 < id2
        MATCH (e1 {id: id1})-[r]-(e2 {id: id2})
        RETURN 
            e1.name AS source,
            type(r) AS relationship,
            e2.name AS target,
            properties(r) AS properties
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(query, ids=entity_ids)
            return [record.data() for record in result]
    
    # =========================================================
    # QUERY AND RESPONSE
    # =========================================================
    
    def query(
        self,
        question: str,
        llm_function: Callable[[str], str] = None,
        **kwargs
    ) -> RAGResponse:
        """
        Execute a RAG query.
        
        Args:
            question: User question
            llm_function: Function to call LLM (takes prompt, returns response)
            **kwargs: Additional query parameters
            
        Returns:
            RAG response with answer and context
        """
        import time
        start = time.perf_counter()
        
        # Create query object
        rag_query = RAGQuery(
            question=question,
            top_k=kwargs.get('top_k', 5),
            include_graph_context=kwargs.get('include_graph_context', True),
            max_context_tokens=kwargs.get('max_context_tokens', 4000),
            min_similarity=kwargs.get('min_similarity', 0.3),
            context_hops=kwargs.get('context_hops', 1)
        )
        
        # Retrieve context
        context = self.hybrid_retrieve(rag_query)
        
        # Build prompt
        prompt = self._build_prompt(question, context)
        
        # Generate answer
        if llm_function:
            answer = llm_function(prompt)
        else:
            answer = f"[Context retrieved - {len(context.chunks)} chunks found. " \
                     f"LLM function not provided for generation.]"
        
        # Calculate processing time
        elapsed_ms = (time.perf_counter() - start) * 1000
        
        # Extract sources
        sources = [c.get('id', 'unknown') for c in context.chunks]
        
        return RAGResponse(
            question=question,
            answer=answer,
            context=context,
            sources=sources,
            confidence=self._calculate_confidence(context),
            processing_time_ms=elapsed_ms
        )
    
    def _build_prompt(
        self,
        question: str,
        context: RetrievedContext
    ) -> str:
        """Build LLM prompt with context."""
        prompt = f"""You are a helpful assistant with access to a knowledge base.
Use the following context to answer the question. If the context doesn't contain
relevant information, say so.

CONTEXT:
{context.to_prompt_context()}

QUESTION: {question}

INSTRUCTIONS:
1. Answer based on the provided context
2. Cite specific sources when possible
3. If uncertain, acknowledge limitations
4. Be concise but thorough

ANSWER:"""
        return prompt
    
    def _calculate_confidence(self, context: RetrievedContext) -> float:
        """Calculate confidence score based on retrieval quality."""
        if not context.chunks:
            return 0.0
        
        # Average similarity score
        scores = [c.get('score', 0) for c in context.chunks]
        avg_score = sum(scores) / len(scores) if scores else 0
        
        # Bonus for having entities and graph context
        entity_bonus = min(len(context.entities) * 0.05, 0.2)
        graph_bonus = 0.1 if context.graph_context else 0
        
        confidence = min(avg_score + entity_bonus + graph_bonus, 1.0)
        return round(confidence, 2)


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Graph RAG Fundamentals - Demonstration")
    print("=" * 60)
    
    # Architecture explanation
    print("""
    Graph RAG combines two powerful retrieval methods:
    
    1. VECTOR SEARCH (Semantic)
       - Finds chunks similar in meaning to the query
       - Uses HNSW index for fast approximate search
       - Returns relevance scores (0-1)
    
    2. GRAPH TRAVERSAL (Structural)
       - Follows relationships from retrieved chunks
       - Adds entity information (Patient, Disease, etc.)
       - Provides relationship context
    
    The combination provides:
    - Semantic relevance (what's the content about?)
    - Structural context (how are things connected?)
    - Better factual grounding for LLM responses
    """)
    
    print("\n📌 Initializing Graph RAG")
    print("-" * 40)
    
    try:
        driver = get_driver()
        
        # Try to use local embeddings
        try:
            embeddings = LocalEmbeddings()
            print("   Using local embeddings (sentence-transformers)")
        except ImportError:
            print("   sentence-transformers not available")
            embeddings = None
        
        if embeddings:
            rag = GraphRAG(
                driver=driver,
                embedding_provider=embeddings,
                vector_index_name="chunk_embeddings"
            )
            
            # Test query (without LLM)
            print("\n📌 Testing RAG Retrieval")
            print("-" * 40)
            
            response = rag.query(
                question="What are the treatment options for diabetes?",
                top_k=3
            )
            
            print(f"\n   Question: {response.question}")
            print(f"   Chunks retrieved: {len(response.context.chunks)}")
            print(f"   Entities found: {len(response.context.entities)}")
            print(f"   Confidence: {response.confidence}")
            print(f"   Processing time: {response.processing_time_ms:.2f}ms")
            
            if response.context.chunks:
                print("\n   Top retrieved chunks:")
                for i, chunk in enumerate(response.context.chunks[:3], 1):
                    content = chunk.get('content', '')[:100]
                    score = chunk.get('score', 0)
                    print(f"   {i}. [{score:.3f}] {content}...")
        
        driver.close()
        
    except Exception as e:
        print(f"   Error: {e}")
    
    print("\n" + "=" * 60)
    print("✅ Graph RAG demonstration complete!")
    print("=" * 60)
